/*
Copyright (C) 2006 Pedro Felzenszwalb

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307 USA
*/

/* basic image I/O */

#pragma once

#include <cstdlib>
#include <climits>
#include <cstring>
#include <fstream>

#include "image.h"
#include "misc.h"

namespace distance_transform
{

  const short BUF_SIZE = 256;

  class pnm_error { };

  static void read_packed(unsigned char *data, int size, std::ifstream &f) {
    unsigned char c = 0;
    
    int bitshift = -1;
    for (int pos = 0; pos < size; pos++) {
      if (bitshift == -1) {
        c = f.get();
        bitshift = 7;
      }
      data[pos] = (c >> bitshift) & 1;
      bitshift--;
      }
  }

  static void write_packed(unsigned char *data, int size, std::ofstream &f) {
    unsigned char c = 0;
    
    int bitshift = 7;
    for (int pos = 0; pos < size; pos++) {
        c = c | (data[pos] << bitshift);
        bitshift--;
        if ((bitshift == -1) || (pos == size-1)) {
    f.put(c);
    bitshift = 7;
    c = 0;
        }
    }
  }

  /* read PNM field, skipping comments */ 
  static void pnm_read(std::ifstream &file, char *buf) {
    char doc[BUF_SIZE];
    char c;
    
    file >> c;
    while (c == '#') {
      file.getline(doc, BUF_SIZE);
      file >> c;
    }
    file.putback(c);
    
    file.width(BUF_SIZE);
    file >> buf;
    file.ignore();
  }

  static image<uchar> *loadPBM(const char *name) {
    char buf[BUF_SIZE];
    
    /* read header */
    std::ifstream file(name, std::ios::in | std::ios::binary);
    pnm_read(file, buf);
    if (strncmp(buf, "P4", 2))
      throw pnm_error();
      
    pnm_read(file, buf);
    int width = atoi(buf);
    pnm_read(file, buf);
    int height = atoi(buf);
    
    /* read data */
    image<uchar> *im = new image<uchar>(width, height);
    for (int i = 0; i < height; i++)
      read_packed(imPtr(im, 0, i), width, file);
    
    return im;
  }

  static void savePBM(image<uchar> *im, const char *name) {
    int width = im->width();
    int height = im->height();
    std::ofstream file(name, std::ios::out | std::ios::binary);

    file << "P4\n" << width << " " << height << "\n";
    for (int i = 0; i < height; i++)
      write_packed(imPtr(im, 0, i), width, file);
  }

  static image<uchar> *loadPGM(const char *name) {
    char buf[BUF_SIZE];
    
    /* read header */
    std::ifstream file(name, std::ios::in | std::ios::binary);
    pnm_read(file, buf);
    if (strncmp(buf, "P5", 2))
      throw pnm_error();

    pnm_read(file, buf);
    int width = atoi(buf);
    pnm_read(file, buf);
    int height = atoi(buf);

    pnm_read(file, buf);
    if (atoi(buf) > UCHAR_MAX)
      throw pnm_error();

    /* read data */
    image<uchar> *im = new image<uchar>(width, height);
    file.read((char *)imPtr(im, 0, 0), width * height * sizeof(uchar));

    return im;
  }

  static void savePGM(image<uchar> *im, const char *name) {
    int width = im->width();
    int height = im->height();
    std::ofstream file(name, std::ios::out | std::ios::binary);

    file << "P5\n" << width << " " << height << "\n" << UCHAR_MAX << "\n";
    file.write((char *)imPtr(im, 0, 0), width * height * sizeof(uchar));
  }

  static image<rgb> *loadPPM(const char *name) {
    char buf[BUF_SIZE], doc[BUF_SIZE];
    
    /* read header */
    std::ifstream file(name, std::ios::in | std::ios::binary);
    pnm_read(file, buf);
    if (strncmp(buf, "P6", 2))
      throw pnm_error();

    pnm_read(file, buf);
    int width = atoi(buf);
    pnm_read(file, buf);
    int height = atoi(buf);

    pnm_read(file, buf);
    if (atoi(buf) > UCHAR_MAX)
      throw pnm_error();

    /* read data */
    image<rgb> *im = new image<rgb>(width, height);
    file.read((char *)imPtr(im, 0, 0), width * height * sizeof(rgb));

    return im;
  }

  static void savePPM(image<rgb> *im, const char *name) {
    int width = im->width();
    int height = im->height();
    std::ofstream file(name, std::ios::out | std::ios::binary);

    file << "P6\n" << width << " " << height << "\n" << UCHAR_MAX << "\n";
    file.write((char *)imPtr(im, 0, 0), width * height * sizeof(rgb));
  }

  template <class T>
  void load_image(image<T> **im, const char *name) {
    char buf[BUF_SIZE];
    
    /* read header */
    std::ifstream file(name, std::ios::in | std::ios::binary);
    pnm_read(file, buf);
    if (strncmp(buf, "VLIB", 9))
      throw pnm_error();

    pnm_read(file, buf);
    int width = atoi(buf);
    pnm_read(file, buf);
    int height = atoi(buf);

    /* read data */
    *im = new image<T>(width, height);
    file.read((char *)imPtr((*im), 0, 0), width * height * sizeof(T));
  }

  template <class T>
  void save_image(image<T> *im, const char *name) {
    int width = im->width();
    int height = im->height();
    std::ofstream file(name, std::ios::out | std::ios::binary);

    file << "VLIB\n" << width << " " << height << "\n";
    file.write((char *)imPtr(im, 0, 0), width * height * sizeof(T));
  }

} // namespace
